from vae.vae_model import VAE
import torch
import matplotlib.pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vae = VAE(1, 256, 16, 784).to(device)

vae.load_state_dict(torch.load("./best_model"))


decoder = vae.decoder



sample = torch.normal(0, 1, size=(64, 16), device=device)

outputs = decoder(sample).detach().cpu().numpy().reshape(-1, 28, 28)

fig, ax = plt.subplots(8, 8)

for i in range(8):
    for j in range(8):
        idx = 8 * i + j
        ax[i, j].imshow(outputs[idx])
        ax[i, j].axis('off')

plt.show()